package com.zt.chatglm.executor.aigc;

import com.alibaba.fastjson.JSON;
import com.zt.chatglm.IOpenAiApi;
import com.zt.chatglm.commons.Constants;
import com.zt.chatglm.domain.chat.ChatCompletionRequest;
import com.zt.chatglm.domain.chat.ChatCompletionResponse;
import com.zt.chatglm.domain.chat.ChatCompletionSyncResponse;
import com.zt.chatglm.domain.image.ImageCompletionRequest;
import com.zt.chatglm.domain.image.ImageCompletionResponse;
import com.zt.chatglm.executor.Executor;
import com.zt.chatglm.executor.result.ResultHandler;
import com.zt.chatglm.session.Configuration;
import okhttp3.*;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import org.jetbrains.annotations.Nullable;

import java.util.List;
import java.util.concurrent.CompletableFuture;

/**
 * @author: zt
 * @Description: TODO
 * @create: 2024-05-20 23:36
 */

public class GLMExecutor implements Executor, ResultHandler {


    /**
     * OpenAi 接口
     */
    private final Configuration configuration;
    /**
     * 工厂事件
     */
    private final EventSource.Factory factory;
    /**
     * 统一接口
     */
    private IOpenAiApi openAiApi;

    private OkHttpClient okHttpClient;

    public GLMExecutor(Configuration configuration) {
        this.configuration = configuration;
        this.factory = configuration.createRequestFactory();
        this.openAiApi = configuration.getOpenAiApi();
        this.okHttpClient = configuration.getOkHttpClient();
    }

    @Override
    public EventSource completions(ChatCompletionRequest completionRequest, EventSourceListener eventSourceListener) throws Exception {
        // 1. 构建请求
        Request request = new Request.Builder()
                // 构建请求url https://open.bigmodel.cn/ + api/paas/v4/chat/completions
                .url(configuration.getApiHost().concat(IOpenAiApi.v4_completions))
                // 构建请求体
                .post(RequestBody.create(MediaType.parse(Configuration.JSON_CONTENT_TYPE), completionRequest.toString()))
                .build();
        // 2. 返回事件结果（兼容处理）
        return factory.newEventSource(request, completionRequest.getIsCompatible() ? eventSourceListener(eventSourceListener) : eventSourceListener);
    }

    @Override
    public CompletableFuture<String> completions(ChatCompletionRequest completionRequest) throws Exception {
        // 1. 构建请求
        Request request = new Request.Builder()
                .url(configuration.getApiHost().concat(IOpenAiApi.v4_completions))
                .post(RequestBody.create(MediaType.parse(Configuration.JSON_CONTENT_TYPE), completionRequest.toString()))
                .build();
        // 2. 包装返回结果消息
        CompletableFuture<String> future = new CompletableFuture<>();
        StringBuffer dataBuffer = new StringBuffer();
        // 3. 异步发送请求
        factory.newEventSource(request, new EventSourceListener() {
            @Override
            public void onEvent(EventSource eventSource, @Nullable String id, @Nullable String type, String data) {
                if ("[DONE]".equals(data)) {
                    return;
                }
                ChatCompletionResponse response = JSON.parseObject(data, ChatCompletionResponse.class);
                List<ChatCompletionResponse.Choice> choices = response.getChoices();
                for (ChatCompletionResponse.Choice choice : choices) {
                    if (!"stop".equals(choice.getFinishReason())) {
                        dataBuffer.append(choice.getDelta().getContent());
                    }
                }
            }

            @Override
            public void onClosed(EventSource eventSource) {
                future.complete(dataBuffer.toString());
            }

            @Override
            public void onFailure(EventSource eventSource, @Nullable Throwable t, @Nullable Response response) {
                future.completeExceptionally(new RuntimeException("Request closed before completion"));
            }
        });
        return future;
    }

    @Override
    public ChatCompletionSyncResponse completionsSync(ChatCompletionRequest completionRequest) throws Exception {
        // 1. 修改请求参数请求类型为同步请求
        completionRequest.setStream(false);
        // 2. 构建请求消息
        Request request = new Request.Builder()
                .url(configuration.getApiHost().concat(IOpenAiApi.v4_completions))
                .post(RequestBody.create(MediaType.parse(Configuration.JSON_CONTENT_TYPE), completionRequest.toString()))
                .build();
        // 3. 使用okhttp发送请求
        OkHttpClient okHttpClient = configuration.getOkHttpClient();
        Response response = okHttpClient.newCall(request).execute();
        if (!response.isSuccessful()) {
            throw new RuntimeException("Request failed");
        }
        return JSON.parseObject(response.body().string(), ChatCompletionSyncResponse.class);
    }

    /**
     * 图片生成
     *
     * Request request = new Request.Builder()
     *         .url(configuration.getApiHost().concat(IOpenAiApi.cogview3))
     *         .post(RequestBody.create(MediaType.parse(Configuration.JSON_CONTENT_TYPE), imageCompletionRequest.toString()))
     *         .build();
     * // 请求结果
     * Call call = okHttpClient.newCall(request);
     * Response execute = call.execute();
     * ResponseBody body = execute.body();
     * if (execute.isSuccessful() && body != null) {
     *     String responseBody = body.string();
     *     ObjectMapper objectMapper = new ObjectMapper();
     *     return objectMapper.readValue(responseBody, ImageCompletionResponse.class);
     * } else {
     *     throw new IOException("Failed to get image response");
     * }
     * @param request 请求信息
     * @return
     * @throws Exception
     */
    @Override
    public ImageCompletionResponse genImages(ImageCompletionRequest request) throws Exception {
        return openAiApi.genImages(request).blockingGet();
    }

    /**
     * 使用装饰者模式对动态的对事件监听源新增一些额外的功能
     * @param eventSourceListener
     * @return
     */
    @Override
    public EventSourceListener eventSourceListener(EventSourceListener eventSourceListener) {
        return new EventSourceListener() {
            @Override
            public void onEvent(EventSource eventSource, @Nullable String id, @Nullable String type, String data) {
                // 1. 结束接受新信息
                if ("[DONE]".equals(data)) {
                    return;
                }
                ChatCompletionResponse response = JSON.parseObject(data, ChatCompletionResponse.class);
                // 2. 新响应的结束标志为stop，这里兼容处理
                if (response.getChoices() != null && 1 == response.getChoices().size() && "stop".equals(response.getChoices().get(0).getFinishReason())) {
                    eventSourceListener.onEvent(eventSource, id, Constants.EventType.finish.getCode(), data);
                    return;
                }
                // 3. 说明消息还未发送结束，继续发送
                eventSourceListener.onEvent(eventSource, id, Constants.EventType.add.getCode(), data);
            }

            @Override
            public void onClosed(EventSource eventSource) {
                eventSourceListener.onClosed(eventSource);
            }

            @Override
            public void onFailure(EventSource eventSource, @Nullable Throwable t, @Nullable Response response) {
                eventSourceListener.onFailure(eventSource, t, response);
            }
        };
    }
}
